//
// Created by craipy on 23-2-7.
//

#ifndef KALMAN_LINEARKALMANFILTER_HPP
#define KALMAN_LINEARKALMANFILTER_HPP

#include "kalman/MeasurementModel.hpp"
#include "kalman/SystemModel.hpp"
#include "kalman/KalmanFilterBase.hpp"
#include "kalman/StandardFilterBase.hpp"
#include "kalman/LinearMeasurementModel.hpp"
#include "kalman/LinearSystemModel.hpp"

namespace Kalman {
    template<class StateType>
    class LinearKalmanFilter : public KalmanFilterBase<StateType>,
                               public StandardFilterBase<StateType> {
    public:
        //! Kalman Filter base type
        typedef KalmanFilterBase<StateType> KalmanBase;
        //! Standard Filter base type
        typedef StandardFilterBase<StateType> StandardBase;

        //! Numeric Scalar Type inherited from base
        using typename KalmanBase::T;

        //! State Type inherited from base
        using typename KalmanBase::State;

        //! Linearized Measurement Model Type
        template<class Measurement>
        using MeasurementModelType = LinearMeasurementModel<State, Measurement>;

        //! Linearized System Model Type
        template<class Control>
        using SystemModelType = LinearSystemModel<State, Control>;

    protected:
        //! Kalman Gain Matrix Type
        template<class Measurement>
        using KalmanGain = Kalman::KalmanGain<State, Measurement>;

    protected:
        //! State Estimate
        using KalmanBase::x;
        //! State Covariance Matrix
        using StandardBase::P;
    public:
        LinearKalmanFilter() {
            P.setIdentity();
        }

        template<class Control>
        const State &predict(SystemModelType<Control> &s) {
            // predict state (without control)
            Control u;
            u.setZero();
            return predict(s, u);
        }

        template<class Control>
        const State &predict(SystemModelType<Control> &s, const Control &u) {
            // predict state
            x = s.f(x, u);

            // predict covariance
            P = (s.A * P * s.A.transpose()) + s.getCovariance();

            // return state prediction
            return x;
        }

        template<class Measurement>
        const State &update(MeasurementModelType<Measurement> &m, const Measurement &z) {
            // COMPUTE KALMAN GAIN
            // compute innovation covariance
            Covariance<Measurement> S = (m.H * P * m.H.transpose()) + (m.getCovariance());

            // compute kalman gain
            KalmanGain<Measurement> K = P * m.H.transpose() * S.inverse();

            // UPDATE STATE ESTIMATE AND COVARIANCE
            // Update state using computed kalman gain and innovation
            x += K * (z - m.h(x));

            // Update covariance
            P -= K * m.H * P;

            // return updated state estimate
            return x;
        }
    };

} // Kalman

#endif //KALMAN_LINEARKALMANFILTER_HPP
